% perform_allocation.m
% 核心内存分配函数，执行地址分配和SPILL决策。
% *** 已修改: 修正超时检测逻辑中的bug ***
function [solution, graph] = perform_allocation(initial_schedule, graph, code_map, l1_cap, ub_cap, l0a_cap, l0b_cap, l0c_cap, min_farthest_use_threshold, w_use, w_vol, w_vic, allocation_strategy, debug_target_step, debug_window)
    % --- 1. 初始化 ---
    solution = struct('final_schedule', [], 'memory_offsets', containers.Map('KeyType','int32','ValueType','any'), 'spill_map', containers.Map('KeyType','int32','ValueType','any'), 'total_spill_volume', 0);
    solution.timed_out = false;
    solution.nodes_processed = 0;
    solution.spill_failed = false;
    
    node_template = struct(...
        'id', [], 'op', '', 'buf_id', [], 'size', [], ...
        'type', '', 'pipe', '', 'cycles', [], 'bufs', [], ...
        'op_code', 0, 'type_code', 0);
    next_node_id = graph.num_nodes + 1;
    
    cache_managers = containers.Map('KeyType', 'int32', 'ValueType', 'any');
    cache_managers(code_map.L1) = cm_create(l1_cap);
    cache_managers(code_map.UB) = cm_create(ub_cap);
    cache_managers(code_map.L0A) = cm_create(l0a_cap);
    cache_managers(code_map.L0B) = cm_create(l0b_cap);
    cache_managers(code_map.L0C) = cm_create(l0c_cap);
    
    active_buffers = containers.Map('KeyType', 'int32', 'ValueType', 'any');
    spilled_buffers = containers.Map('KeyType', 'int32', 'ValueType', 'any');
    copy_in_bufs = [];
    for i = 1:graph.num_nodes
        if isempty(graph.nodes(i).id)
            continue;
        end
        node = graph.nodes(i);
        if strcmp(node.op, 'COPY_IN') && ~isempty(node.bufs)
            copy_in_bufs = [copy_in_bufs, node.bufs];
        end
    end
    copy_in_bufs = unique(copy_in_bufs);
    
    buf_id_to_free_step = containers.Map('KeyType', 'int32', 'ValueType', 'any');
    for i = 1:length(initial_schedule)
        node_id = initial_schedule(i);
        if node_id <= length(graph.nodes) && ~isempty(graph.nodes(node_id).id)
            node = graph.nodes(node_id);
            if node.op_code == code_map.FREE && ~isempty(node.buf_id)
                buf_id_to_free_step(node.buf_id) = i;
            end
        end
    end
    
    % --- 2. 主循环 ---
    step = 1;
    schedule = initial_schedule;
    max_steps = length(initial_schedule) * 3;
    
    while step <= length(schedule)
        if step > max_steps
            processed_nodes_in_current_schedule = schedule(1:step-1);
            is_initial_node = ismember(processed_nodes_in_current_schedule, initial_schedule);
            solution.nodes_processed = sum(is_initial_node);
            solution.timed_out = true;
            % *** 已修改: 删除此处的错误代码行 ***
            % 此处无需任何操作，solution.total_spill_volume 已是正确的累加值
            break; 
        end
        node_id = schedule(step);
        node = graph.nodes(node_id);
        is_alloc_request = (node.op_code == code_map.ALLOC);
        is_spill_in_request = strcmp(node.op, 'SPILL_IN');
        if is_alloc_request || is_spill_in_request
            buf_id_to_alloc = node.buf_id;
            if isempty(buf_id_to_alloc)
                step = step + 1; continue;
            end
            
            if is_alloc_request
                if isempty(node.size) || node.size <= 0 || node.type_code == 0
                    step = step + 1; continue;
                end
                size_needed = node.size;
                type_code = node.type_code;
                if isKey(active_buffers, buf_id_to_alloc)
                    error('Allocator 逻辑错误: 在步骤 %d, 节点 %d, 尝试重复分配一个仍在使用的 BufID: %d', step, node_id, buf_id_to_alloc);
                end
            else % is_spill_in_request
                if ~isKey(spilled_buffers, buf_id_to_alloc)
                    step = step + 1; continue;
                end
                spill_info = spilled_buffers(buf_id_to_alloc);
                type_code = spill_info(1);
                size_needed = spill_info(2);
            end
            manager = cache_managers(type_code);
            
            current_buf_free_step = -1;
            if isKey(buf_id_to_free_step, buf_id_to_alloc)
                current_buf_free_step = buf_id_to_free_step(buf_id_to_alloc);
            end
            [manager, offset] = cm_allocate(manager, size_needed, buf_id_to_alloc, allocation_strategy, buf_id_to_free_step, current_buf_free_step);
            
            cache_managers(type_code) = manager;
            if offset ~= -1
                if is_alloc_request
                    solution.memory_offsets(buf_id_to_alloc) = offset;
                else % is_spill_in_request
                    solution.spill_map(node_id) = offset;
                    remove(spilled_buffers, buf_id_to_alloc);
                end
                active_buffers(buf_id_to_alloc) = [type_code, offset, size_needed];
                step = step + 1;
                continue;
            
            else % Spill logic
                [best_plan, success] = select_spill_plan(graph, active_buffers, manager, size_needed, type_code, schedule, step, min_farthest_use_threshold, w_use, w_vol, w_vic);
                
                if ~success
                    % --- BEGIN GRACEFUL FAILURE HANDLING ---
                    % Spill失败，不中断执行，而是设置失败标志并返回，以便上层施加惩罚。
                    solution.spill_failed = true;
                    solution.nodes_processed = step; % 记录失败时的进度 

                    return; % 立即退出函数
                    % --- END GRACEFUL FAILURE HANDLING ---
                end

                insertions = [];
                for victim_buf_id = best_plan.victim_ids
                    victim_info = active_buffers(victim_buf_id);
                    victim_type_code = victim_info(1);
                    victim_size = victim_info(3);
                    
                    if ismember(victim_buf_id, copy_in_bufs)
                        solution.total_spill_volume = solution.total_spill_volume + victim_size;
                    else
                        solution.total_spill_volume = solution.total_spill_volume + victim_size * 2;
                    end
                    
                    spilled_buffers(victim_buf_id) = [victim_type_code, victim_size];
                    remove(active_buffers, victim_buf_id);
                    
                    manager = cache_managers(victim_type_code);
                    manager = cm_free(manager, victim_buf_id);
                    cache_managers(victim_type_code) = manager;
                    spill_out_node = node_template;
                    spill_out_node.id = next_node_id;
                    spill_out_node.op = 'SPILL_OUT';
                    spill_out_node.bufs = victim_buf_id;
                    spill_out_node.buf_id = victim_buf_id;
                    
                    spill_in_node = node_template;
                    spill_in_node.id = next_node_id + 1;
                    spill_in_node.op = 'SPILL_IN';
                    spill_in_node.bufs = victim_buf_id;
                    spill_in_node.buf_id = victim_buf_id;
                    spill_out_node.pipe = 'MTE3';
                    spill_in_node.pipe = 'MTE2';
                    if ismember(victim_buf_id, copy_in_bufs)
                        spill_out_node.cycles = 0;
                        spill_in_node.cycles = victim_size * 2 + 150;
                    else
                        spill_out_node.cycles = victim_size * 2 + 150;
                        spill_in_node.cycles = victim_size * 2 + 150;
                    end
                    
                    next_node_id = next_node_id + 2;
                    graph.nodes(spill_out_node.id) = spill_out_node;
                    graph.nodes(spill_in_node.id) = spill_in_node;
                    graph.num_nodes = next_node_id - 1;
                    new_size = graph.num_nodes;
                    if new_size > length(graph.adj_list)
                        graph.adj_list{new_size} = []; 
                        graph.predecessors_list{new_size} = [];
                    end
                    
                    victim_detail = best_plan.victim_details(victim_buf_id);
                    farthest_use = victim_detail{1};
                    insertion_point = farthest_use;
                    if insertion_point > length(schedule), insertion_point = length(schedule) + 1; end
                    
                    insertions = [insertions; step, spill_out_node.id];
                    insertions = [insertions; insertion_point, spill_in_node.id];
                end
                
                insertions = sortrows(insertions, 1, 'descend');
                for ins_idx = 1:size(insertions, 1)
                    point = insertions(ins_idx, 1);
                    id_to_insert = insertions(ins_idx, 2);
                    schedule = [schedule(1:point-1); id_to_insert; schedule(point:end)];
                end
            end
        elseif node.op_code == code_map.FREE
            if ~isempty(node.buf_id) && isKey(active_buffers, node.buf_id)
                buf_id = node.buf_id;
                buffer_info = active_buffers(buf_id);
                type_code = buffer_info(1);
                
                manager = cache_managers(type_code);
                manager = cm_free(manager, buf_id);
                cache_managers(type_code) = manager;
                
                remove(active_buffers, buf_id);
            end
            step = step + 1;
        else
            step = step + 1;
        end
    end
    solution.final_schedule = schedule;
end